use std::{convert::Infallible, net::SocketAddr};

use axum::{
    extract::{
        ws::{Message, WebSocket},
        ConnectInfo, WebSocketUpgrade,
    },
    response::{sse::Event, IntoResponse, Response, Sse},
};
use axum_extra::{headers, TypedHeader};
use futures::{stream, Stream, StreamExt};
use tokio::time::{sleep, Duration};

pub async fn sokcet_handler(
    ws: WebSocketUpgrade,
    _sec_websocket_key: Option<TypedHeader<headers::SecWebsocketKey>>, // 获取请求头地址
    addr: ConnectInfo<SocketAddr>, // 需要使用into_make_service_with_connect_info
) -> Response {
    println!("addr: {:?}", addr);
    ws.on_upgrade(handler)
}
async fn handler(mut socket: WebSocket) {
    while let Some(Ok(msg)) = socket.recv().await {
        let is_err = match msg {
            Message::Text(msg) => socket.send(Message::Text(msg)).await.is_err(),
            Message::Binary(bin) => socket.send(Message::Binary(bin)).await.is_err(),
            Message::Ping(msg) => socket.send(Message::Ping(msg)).await.is_err(),
            Message::Pong(msg) => socket.send(Message::Pong(msg)).await.is_err(),
            Message::Close(_) => true,
        };
        if true == is_err {
            break;
        }
    }
    println!("发生异常，连接关闭！");
}

// eventSource 依赖 tokio_stream, futures
pub async fn sse_handler() -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    let mut i = 0;
    let stream = stream::repeat_with(move || {
        i += 1;
        Event::default().data(format!("hi, {}", i))
    })
    .map(Ok);

    // 每隔1秒发1次保活
    Sse::new(stream).keep_alive(
        axum::response::sse::KeepAlive::new()
            .interval(Duration::from_secs(1))
            .text("keep-alive-text"),
    )
}

// 逐步响应
pub async fn post_step() -> impl IntoResponse {
    let stream = stream::iter(0..10).map(async move |i| {
        sleep(Duration::from_millis(300)).await;
        let event = Event::default().data(format!("message: {},", i));
        Ok::<_, std::io::Error>(event)
    });
    
    Sse::new(stream.buffer_unordered(1))
}